import pandas as pd
import re

def merge_excel_files(file_a_path, file_b_path, output_path):
    # Read both Excel files
    df_a = pd.read_excel(file_a_path)
    df_b = pd.read_excel(file_b_path)
    
    # Create a new dataframe
    result_df = pd.DataFrame()
    
    # Add 编号 column first
    result_df["编号"] = range(1, len(df_a) + 1)
    
    # Then add 中文名 and 英文名 columns
    if "中文名" in df_a.columns:
        result_df["中文名"] = df_a["中文名"]
    else:
        result_df["中文名"] = None
        
    if "英文名" in df_a.columns:
        result_df["英文名"] = df_a["英文名"]
    else:
        result_df["英文名"] = None
    
    # List of required columns we want to ensure exist
    required_columns = ["CID号", "CAS号", "分子式", "SMILES", "二级碎片", "活性", "毒性", 
                       "药理作用", "参考文献", "类别", "CB号", "CB网站", "PC网站", "本地数据库", "CB使用中文名"]
    
    # Add existing required columns from df_a if they exist
    for col in required_columns:
        if col in df_a.columns:
            result_df[col] = df_a[col]
        else:
            result_df[col] = None
    
    # Process each row in result_df
    for idx, row in result_df.iterrows():
        # Skip processing if both 中文名 and 英文名 are NaN
        if pd.isna(row.get("中文名")) and pd.isna(row.get("英文名")):
            continue
            
        # Try to find matches in df_b
        matches_zh = df_b[df_b["NameZH"] == row["中文名"]] if not pd.isna(row.get("中文名")) else pd.DataFrame()
        matches_en = df_b[df_b["NameEN"] == row["英文名"]] if not pd.isna(row.get("英文名")) else pd.DataFrame()
        
        # Combine matches
        matches = pd.concat([matches_zh, matches_en]).drop_duplicates()
        
        if not matches.empty:
            # Get the first match (if multiple)
            match = matches.iloc[0]
            result_df.at[idx, "本地数据库"] = "Yes"
            # Fill missing Chinese name if available
            if pd.isna(row["中文名"]) and not pd.isna(match["NameZH"]):
                result_df.at[idx, "中文名"] = match["NameZH"]
            
            # Fill missing English name if available
            if pd.isna(row["英文名"]) and not pd.isna(match["NameEN"]):
                result_df.at[idx, "英文名"] = match["NameEN"]
            
            # Insert CID if it doesn't start with "Miss"
            if "CID" in match and not pd.isna(match["CID"]):
                if not str(match["CID"]).startswith("Miss"):
                    result_df.at[idx, "CID号"] = match["CID"]
            
            # Insert CAS if it only contains numbers and dashes
            if "CAS" in match and not pd.isna(match["CAS"]):
                cas_value = str(match["CAS"])
                if re.match(r'^[0-9-]+$', cas_value):
                    result_df.at[idx, "CAS号"] = match["CAS"]
            
            # Insert SMILES
            if "Smiles" in match and not pd.isna(match["Smiles"]):
                result_df.at[idx, "SMILES"] = match["Smiles"]
                
            if "MolecularFormula" in match and not pd.isna(match["MolecularFormula"]):
                result_df.at[idx, "分子式"] = match["MolecularFormula"]
                
            if "MS_Fragment" in match and not pd.isna(match["MS_Fragment"]):
                result_df.at[idx, "二级碎片"] = match["MS_Fragment"]
                
            if "Activity_IC50" in match and not pd.isna(match["Activity_IC50"]):
                result_df.at[idx, "活性"] = match["Activity_IC50"]
                
            if "Virulence_IC50" in match and not pd.isna(match["Virulence_IC50"]):
                result_df.at[idx, "毒性"] = match["Virulence_IC50"]
                
            if "Pharmacology" in match and not pd.isna(match["Pharmacology"]):
                result_df.at[idx, "药理作用"] = match["Pharmacology"]
            
            if "Refer" in match and not pd.isna(match["Refer"]):
                result_df.at[idx, "参考文献"] = match["Refer"]
    
    # Save the result to a new Excel file
    result_df.to_excel(output_path, index=False)
    print(f"Merged file saved as {output_path}")

# Example usage
if __name__ == "__main__":
    file_a_path = "./API/xbzk1.xlsx"
    file_b_path = "../../../chem-app/chemapp_resource/pyscript/CreateDB/output_AllTable.xlsx"
    output_path = "./API/xbzk1_Merged.xlsx"
    
    merge_excel_files(file_a_path, file_b_path, output_path)